"""
IDM.py
"""
import math
import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import multivariate_normal
from driver_model.Moveable import Moveable
import copy
# from driver_model.Moveable import Moveable

import sys
sys.path.append('..')
# from WOMD_Dataset.Dataloader.src import waymo_training_math_tools as math_tool
# from WOMD_Dataset.Dataloader.src import waymo_training_plot_tools as plot_tool
# from WOMD_Dataset.Dataloader.src import waymo_scenario_classes as class_tool
# from WOMD_Dataset.Dataloader.src import waymo_transfer_training_to_tf as transfer_tool
import tensorflow
# 文件路径
# data_set='/home/mdmp/prpproject/WOMD_Dataset/training_dataset/uncompressed_scenario_training_training.tfrecord-00510-of-01000'
# # 多文件读取
# waymo_dataset_list = []
#
# waymo_dataset_list.append(data_set)
# # 读取数据集
# scenarios_get_list = [0]
# scenarios_summary= transfer_tool.process_waymo_training_dataset(data_set, scenarios_get_list)
#
# # 读取scenario
# scenario = scenarios_summary.scenarios_list[0]
#
# agents_list = scenario.agents_list

'''
Basis classes for the microscopic traffic model IDM (intelligent-driver model,
 * see <a href="http://xxx.uni-augsburg.de/abs/cond-mat/0002177"> M. Treiber, A.
 * Hennecke, and D. Helbing, Congested Traffic States in Empirical Observations
 * and Microscopic Simulations, Phys. Rev. E 62, 1805 (2000)].</a> <br>
 * <br>

IDM:
    general IDM follower strategy realization

'''
ismax = 100 # ve(s) = ve(ismax) if s > ismax
MAX_BRAKING = 10 # m/s^2
DELTA = 4 # Acceleration exponent
MAX_STEER = 0.15


# Pure Pursuit计算
def cal_target_index(robot_state, refer_pathX,refer_pathY, l_d=2.0):
    """得到前视目标点
    Args:
        robot_state (_type_): 当前车辆位置
        refer_path (_type_): 参考轨迹（数组）
        l_d：前视距离，先尝试2.0m
    Returns:
        _type_: 前视目标点的索引
    """
    dists = []
    refer_path = list(zip(refer_pathX,refer_pathY))
    # print(refer_path)
    for xy in refer_path:
        xy = np.array(xy)
        dis = np.linalg.norm(robot_state-xy)
        dists.append(dis)

    min_index = np.argmin(dists)
    # print(refer_path)
    delta_l = np.linalg.norm(np.array(refer_path[min_index])-np.array(robot_state))
    # 搜索前视目标点
    while l_d > delta_l and (min_index+1) < len(refer_path):
        delta_l = np.linalg.norm(np.array(refer_path[min_index+1])-np.array(robot_state))
        min_index += 1
    return min_index

def pure_pursuit_control(robot_state,current_ref_point,width,theta0,l_d=2.0, max_steer=0.15):
    """pure pursuit

    Args:
        robot_state (_type_): 车辆位置
        current_ref_point (_type_): 当前参考路点
        l_d：前视距离
    return:返回前轮转向角delta
    """
    alpha = math.atan2(current_ref_point[1]-robot_state[1], current_ref_point[0]-robot_state[0])-theta0 # 注意此处的theta0的正负
    delta = math.atan2(2*width*np.sin(alpha),l_d)
    if max_steer > abs(delta):
        return delta
    elif delta >= 0:
        return max_steer
    else:
        return -max_steer
    


class IDM():
    '''
    The IDM is a static class for calculation based on condition of specific car.
    It is shared between all the cars with the same
    
    parameters:
    ismax； max distance effect for the car following model.
    MAX_BRAKING: max deceleration m/s^2
    '''

    # ismax specify the distance where
    ismax = 100 # ve(s) = ve(ismax) if s > ismax
    MAX_BRAKING = 10 # m/s^2
    DELTA = 4 # Acceleration exponent
    def __init__(self, v0, a, b,sigma,delta=0.1,s0=2,T=1.5,size=20, max_steer = 0.15):
        global MAX_STEER
        '''
        :param delta:
            technical term in the position calculation, s
        :param v0:
            max speed of the car(desired speed)
        :param a:
            max acceleration m/s^2
        :param b:
            normal deceleration m/s^2
        :param s0: default 2
            least safe distance between two cars
            --> exceed this means need to brake immediately
        :param T:  default 1.5
            human reaction time(s) T for s*
        '''
        global DELTA
        global ismax
        global MAX_BRAKING
        self.v0, self.Delta, self.a, self.b, self.s0,self.T = v0, DELTA, a, b, s0, T
        self.size,self.sigma = size,sigma
        self.delta = delta
        self.sqrtab = np.sqrt(a*b)
        self.veq_table = np.zeros(self.ismax + 1)   # equilibrium velocity table
        self.initialize() #generate equilibrium velocity table
        self.max_steer = max_steer
        MAX_STEER = max_steer


    def initialize(self):
        """
        计算在各种距离下的均衡速度
        """
        dt = 0.5 #relaxation timestep 0.5s
        kmax = 20 #number of iteration in relaxation（迭代次数）
        for s in range(1, self.ismax+1):    # 以1到ismax为不同距离，来生成均衡速度表
            Ve = self.veq_table[s-1]
            for k in range(0, kmax):
                s_star = self.s0 + Ve*self.T
                acc = self.a * (1.-np.power(Ve/self.v0, self.Delta) - (s_star**2) / (s**2))
                Ve += acc * dt
                Ve = max(Ve, 0) # can't be lower
            self.veq_table[s] = Ve

    def Veq(self, dx):
        '''
        function for equilibrium velocity using veq_table;
            ve(s>ismax)=ve(ismax)
        this value is used to set up initial speed of vehicle
        :param dx: the distance between the current and forward car
        :return: velocity in m/s
        '''
        s = int(np.floor(dx))
        V = 0
        if s < 0:
            pass # V=0
        elif s < self.ismax:
            rest = dx - s
            V = (1-rest) * self.veq_table[s] + rest * self.veq_table[s+1]
        else:
            V = self.veq_table[self.ismax]
        return V



    def calc_acc(self, bwd: Moveable, fwd: Moveable):
        '''
        Moveable是车辆的基类，包含了车辆的位置、速度、加速度等信息
        :param bwd: Moveable, The current vehicle
        :param fwd: Moveable, The vehicle in the forward vehicle
        :return: acceleration m/s^2
        '''
        # print(self.veq_table)
        if(fwd == None):
            return 0
        delta_v = bwd.vel[2] - fwd.vel[2]
        s = bwd.distance_to(fwd)
        vel = bwd.vel[2]
        s_star_raw = self.s0 + vel * self.T\
                        + (vel * delta_v) / (2 * self.sqrtab)
        s_star = max(s_star_raw, self.s0)
        acc = self.a * (1 - np.power(vel / self.v0, self.Delta) - (s_star **2) / (s**2))
        acc = max(acc, -self.b)
        # print("acc: ", acc, end=" ")
        # if acc > 0:
        #     acc *= 2
        return acc

    def get_v0(self):
        '''
        get maximum designed speed of this vehicle
        '''
        return self.v0
    
    def get_position(self,bwd:Moveable,fwd:Moveable,pathX,pathY):
        """在舵角方向不变的情况下，做匀加速运动，计算车辆的位置
        增加pure pursuit的舵角变化
        Args:
            bwd (Moveable): _description_
            fwd (Moveable): _description_
        """
        new_vel = [0,0,0]
        pos1= bwd.pos
        theta0 = math.atan2(bwd.vel[1], bwd.vel[0])
        min_index = cal_target_index(bwd.pos,pathX,pathY)
        deltaTheta = pure_pursuit_control(bwd.pos,[pathX[min_index],pathY[min_index]],bwd.width,theta0, max_steer=self.max_steer)
        # 角度的正负要靠尝试试出来
        theta0 = theta0+deltaTheta
        new_pos = [0,0] # new_pos[0] is the x coordinate, new_pos[1] is the y coordinate
        a = self.calc_acc(bwd,fwd)
        # print("a",a)
        # print("theta0",theta0)
        new_pos[0]=pos1[0]+bwd.vel[2]*math.cos(theta0)*self.delta+0.5*a*math.cos(theta0)*self.delta**2
        new_pos[1]=pos1[1]+bwd.vel[2]*math.sin(theta0)*self.delta+0.5*a*math.sin(theta0)*self.delta**2
        new_vel[0] = bwd.vel[0]+a*math.cos(theta0)*self.delta
        new_vel[1] = bwd.vel[1]+a*math.sin(theta0)*self.delta
        new_vel[2] = math.sqrt((bwd.vel[0]+a*math.cos(theta0)*self.delta)**2+(bwd.vel[1]+a*math.sin(theta0)*self.delta)**2)
        new_pos = np.array(new_pos)
        return new_pos,new_vel
    
    def get_position_improve(self,bwd:Moveable,fwd:Moveable,pathX,pathY):
        """在舵角方向不变的情况下，做匀加速运动，计算车辆的位置
        增加pure pursuit的舵角变化
        Args:
            bwd (Moveable): _description_
            fwd (Moveable): _description_
        """
        v1= bwd.vel
        pos1= bwd.pos
        theta = math.atan2(bwd.vel[1], bwd.vel[0])
        look_ahead = 15
        min_index = cal_target_index(bwd.pos,pathX,pathY)
        target_index = min_index + look_ahead
        if target_index >= len(pathX):
            target_index = len(pathX) - 1
        deltaTheta = pure_pursuit_control(bwd.pos,[pathX[target_index],pathY[target_index]],bwd.width,theta, max_steer=self.max_steer)
        # 角度的正负要靠尝试试出来
        theta0 = theta+deltaTheta
        new_pos = [0,0] # new_pos[0] is the x coordinate, new_pos[1] is the y coordinate
        a = self.calc_acc(bwd,fwd)
        # print("a",a)
        # print("theta0",theta0)
        new_pos[0]=pos1[0]+bwd.vel[2]*math.cos(theta0)*self.delta
        new_pos[1]=pos1[1]+bwd.vel[2]*math.sin(theta0)*self.delta
        predictPathX = [(pos1[0]+bwd.vel[0]*t+0.5*a*math.cos(theta0)*t**2) for t in np.arange(0,self.delta+self.delta/self.size,self.delta/self.size)]
        predictPathY = [(pos1[1]+bwd.vel[1]*t+0.5*a*math.sin(theta0)*t**2) for t in np.arange(0,self.delta+self.delta/self.size,self.delta/self.size)]
        predictPath = [predictPathX,predictPathY]
        new_heading = theta + self.delta * math.sin(deltaTheta) * bwd.vel[2] / 2
        new_pos = np.array(new_pos)
        return new_pos,theta0,a, new_heading
    
    def getTrajectory(self,bwd:Moveable,fwd:Moveable,pathX,pathY, points_number = None):
        """
        利用idm和CA模型计算轨迹
        """
        bwd = copy.deepcopy(bwd)
        fwd = copy.deepcopy(fwd)
        if pathX is None:
            return[[0], [0]]

        predictPath = [[],[]]
        originalPos1 = bwd.pos

        if points_number is None:
            points_number = self.size
        
        originalVel1 = bwd.vel
        if not fwd==None:
            originalPos2 = fwd.pos
            originalVel2 = fwd.vel
            for index in range(0, points_number):
                new_pos,theta0,a,new_heading = self.get_position_improve(bwd,fwd,pathX,pathY)
                # fa = math.sqrt(fwd.acc[0]**2+fwd.acc[1]**2)
                thetaFwd = math.atan2(fwd.vel[1], fwd.vel[0])
                min_index = cal_target_index(fwd.pos,pathX,pathY)
                deltaTheta = pure_pursuit_control(fwd.pos,[pathX[min_index],pathY[min_index]],fwd.width,thetaFwd, max_steer=self.max_steer)
                thetaFwd = thetaFwd+deltaTheta
                predictPath[0].append(new_pos[0])
                predictPath[1].append(new_pos[1])
                bwd.acc = [bwd.acc[0]+a*math.cos(theta0),bwd.acc[1]+a*math.sin(theta0)]
                bwd.pos = [new_pos[0],new_pos[1]]
                bwd.vel[2] = bwd.vel[2] + a * self.delta
                if(bwd.vel[2] < 0):
                    bwd.vel[2] = 0.000000001
                bwd.vel[0] = bwd.vel[2] * math.cos(new_heading)
                bwd.vel[1] = bwd.vel[2] * math.sin(new_heading)
                fwd.pos = [fwd.pos[0]+fwd.vel[2]*math.cos(thetaFwd)*self.delta+0.5*fwd.acc*math.cos(thetaFwd)*self.delta**2,fwd.pos[1]+fwd.vel[2]*math.sin(thetaFwd)*self.delta+0.5*fwd.acc*math.sin(thetaFwd)*self.delta**2]
                # fwd.vel = [fwd.vel[0]+fa*math.cos(thetaFwd)*self.delta,fwd.vel[1]+fa*math.sin(thetaFwd)*self.delta,math.sqrt((fwd.vel[0]+fa*math.cos(thetaFwd)*self.delta)**2+(fwd.vel[1]+fa*math.sin(thetaFwd)*self.delta)**2)]
                fwd.vel[2] = fwd.vel[2] + fwd.acc * self.delta
                fwd.vel[0] = fwd.vel[2] * math.cos(thetaFwd)
                fwd.vel[1] = fwd.vel[2] * math.sin(thetaFwd)

            bwd.pos = originalPos1
            fwd.pos = originalPos2
            bwd.vel = originalVel1
            fwd.vel = originalVel2
        else:
            for index in range(0, points_number):
                new_pos,theta0,a, new_heading = self.get_position_improve(bwd,fwd,pathX,pathY)
                predictPath[0].append(new_pos[0])
                predictPath[1].append(new_pos[1])
                bwd.acc = [bwd.acc[0]+a*math.cos(theta0),bwd.acc[1]+a*math.sin(theta0)]
                bwd.pos = [new_pos[0],new_pos[1]]
                bwd.vel[2] = bwd.vel[2] + a * self.delta
                if bwd.vel[2] < 0:
                    bwd.vel[2] = 0.000000001
                bwd.vel[0] = bwd.vel[2] * math.cos(new_heading)
                bwd.vel[1] = bwd.vel[2] * math.sin(new_heading)
            bwd.pos = originalPos1
            bwd.vel = originalVel1
                
        return predictPath


            

    def getPosDistribution(self,bwd: Moveable, fwd: Moveable,pathX,pathY):
        """获取车辆的位置的分布——使用Gaussian Distribution

        Args:
            bwd (Moveable): _description_
            fwd (Moveable): _description_
        """
        next_pos,_ = self.get_position(bwd,fwd,pathX,pathY)
        if not isinstance(next_pos,np.ndarray):
            next_pos = np.array(next_pos)
        # print(next_pos)
        # print("now position is:",bwd.pos)
        cov = np.array([[self.sigma, 0], [0, self.sigma]])
        deltaX = 0.1
        x,y = np.mgrid[bwd.pos[0]-deltaX:bwd.pos[0]+deltaX:.01, bwd.pos[1]-deltaX:bwd.pos[1]+deltaX:.01]  # 做测试使用的x,y，宽度定为±1
        # debug先查测试集
        pos = np.dstack((x, y))
        rv = multivariate_normal(next_pos, cov)
        pos_distribution = rv.pdf(pos)
        return rv.pdf,x,y,pos_distribution



if __name__ == "__main__":
    idm = IDM(13.89,10,4,1,s0=3.0,size=10)
    bwd = Moveable(1,[4933.2101031,651.26111247],[-0.42608935,1.62828386,1.68311035],1,4,4)
    # fwd = Moveable(1,[10,0],[17,0,17],1,4,4)
    fwd = None
    # acc = idm.calc_acc(bwd,fwd)
    # pos,_ = idm.get_position(bwd,fwd,[0,1,2,3,4,5,6,7,8,9,10],[0,0,0,0,0,0,0,0,0,0,0])
    testPathX = [4934.532876379206, 4934.5334571747735, 4934.5332507425155, 4934.531470968702, 4934.527334190091, 4934.520060601764, 4934.5088757522935, 4934.493012152702, 4934.471711016411, 4934.444224154005, 4934.4098160399935, 4934.367766050281, 4934.31737090338, 4934.257947263064, 4934.188834548754, 4934.109397892769, 4934.0190312643, 4933.9171607018825, 4933.803247659658, 4933.67679238804, 4933.5373373421635, 4933.384470541379, 4933.217828838777, 4933.037101041182, 4932.842030805552, 4932.632419265436, 4932.408127312101, 4932.1690774694425, 4931.915255295208, 4931.646710270156, 4931.363556098417, 4931.065970400516, 4930.754193735564, 4930.428527961863, 4930.089333890957, 4929.737028253624, 4929.3720799884295, 4928.995005866046, 4928.606365514183, 4928.206755874889, 4927.796805181536, 4927.377166517701, 4926.948511077003, 4926.511521205945, 4926.066883347526, 4925.615281021889, 4925.157387949884, 4924.693861461093, 4924.225336321316, 4923.752419113116, 4923.275683289867, 4922.7956650660235, 4922.312860237584, 4921.827722106047, 4921.340660590587, 4920.85204269379, 4920.362194430831, 4919.871404356996, 4919.37992882327, 4918.887999114759, 4918.395830589734, 4917.9036339965905]
    testPathY = [641.0788720686771, 641.5756107837692, 642.0723496911489, 642.5690852461012, 643.0658064551221, 643.5624912240276, 644.0591028037805, 644.5555863498009, 645.0518656467673, 645.5478400201921, 646.043381479686, 646.5383321128282, 647.0325017863813, 647.5256661761289, 648.0175651584362, 648.5079016147552, 648.9963406640495, 649.4825093609661, 649.9659968920627, 650.4463552874308, 650.9231006745068, 651.3957150764353, 651.8636487809924, 652.3263232650936, 652.7831346741003, 653.2334578433141, 653.6766508348658, 654.1120599545357, 654.5390252161948, 654.9568861837295, 655.3649881557733, 655.7626886041963, 656.1493638033065, 656.5244155685914, 656.8872780167376, 657.2374242673338, 657.5743729924776, 657.8976947362663, 658.2070179127585, 658.5020344225114, 658.7825047986441, 659.048262845387, 659.2992197092237, 659.5353673526819, 659.7567814142191, 659.9636234510558, 660.1561425736217, 660.33467649841, 660.4996520665244, 660.6515852815058, 660.791080931062, 660.9188318738677, 661.0356180797028, 661.142305497003, 661.239844853425, 661.3292704571999, 661.4116990749295, 661.4883289457165, 661.560438952121, 661.629387961338, 661.6966142845849, 661.7636351821974]
    plt.plot(testPathX,testPathY)
    path = idm.getTrajectory(bwd,fwd,testPathX,testPathY)
    print(path)
    plt.plot(bwd.pos[0], bwd.pos[1], "bo")
    plt.plot(path[0],path[1])
    plt.show()

    # bwd_index = math_tool.Id_to_Index(agents_list,1027)
    # fwd_index = math_tool.Id_to_Index(agents_list,1032)
    # bwd_agent = agents_list[bwd_index]
    # fwd_agent = agents_list[fwd_index]
    # bwd_state = bwd_agent.states_array
    # fwd_state = fwd_agent.states_array
    # bwdState = bwd_state[0,:]
    # fwdState = fwd_state[0,:]
    # pos_array1 = [bwdState[3],bwdState[4]]
    # pos_array2 = [fwdState[3],fwdState[4]]
    # vel_array1 = [bwdState[10],bwdState[11],bwdState[12]]
    # vel_array2 = [fwdState[10],fwdState[11],fwdState[12]]
    # print(vel_array1)
    # print(vel_array2)
    # bwd = Moveable(1,pos_array1,vel_array1,1,bwdState[6],bwdState[7])
    # fwd = Moveable(1,pos_array2,vel_array2,1,fwdState[6],fwdState[7])
    # print("-------------------initial finished-----------------------------")

    # # plot init
    # figure, ax = plt.subplots(figsize=(4,5))
    # plt.ion()
    # plot1, = ax.plot(pos_array1[0], pos_array1[1])
    # plt.xlabel("X-Axis",fontsize=18)
    # plt.ylabel("Y-Axis",fontsize=18)

    # # test demo
    # for time_index in range(1,71):
    #     bwdState = bwd_state[time_index,:]
    #     fwdState = fwd_state[time_index,:]
    #     pos_array1 = [bwdState[3],bwdState[4]]
    #     pos_array2 = [fwdState[3],fwdState[4]]
    #     vel_array1 = [bwdState[10],bwdState[11],bwdState[12]]
    #     vel_array2 = [fwdState[10],fwdState[11],bwdState[12]]

    
